package com.atguigu.userprofile.ml.train

import com.atguigu.userprofile.ml.pipline.MyPipeline
import org.apache.spark.SparkConf
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object StudentGenderTrain {

    def main(args: Array[String]): Unit = {
        //1. 得到数据
        //2. 把数据分成 训练数据集  和 测试数据集 （28原则）
        //3. 初始化流水线对象 填写参数 init
        //4. 把数据送入流水线进行训练


        //1. 得到数据
            // 查询student表  dataframe

        //0  spark运行环境
        val sparkConf: SparkConf = new SparkConf().setAppName("task_ml")
        .setMaster("local[*]")
        val sparkSession: SparkSession = SparkSession.builder().config(sparkConf).enableHiveSupport().getOrCreate()
        //注意：标签值  保持原值就好  但是特征值  必须转换为数字

        val sql =  """
                     |  select uid,
                     |     case hair when '长发' then 110
                     |               when '短发' then 110
                     |               when '板寸' then 110  end as hair,
                     |     height,
                     |     case skirt when '是'   then 11
                     |               when '否'   then 10   end as skirt,
                     |     case age  when '80后' then 80
                     |               when '90后' then 90
                     |               when '00后' then 100  end as age,
                     |     gender
                     |  from student
                     |""".stripMargin
        println(sql)

        println("1. 得到数据")
        sparkSession.sql("use user_profile0318")
        val dataFrame: DataFrame = sparkSession.sql(sql)

        //2. 把数据分成 训练数据集  和 测试数据集 （37原则）
        println("2. 把数据分成 训练数据集  和 测试数据集")
        val Array(trainDF,testDF) = dataFrame.randomSplit(Array(0.7, 0.3))

        //3. 初始化流水线对象 填写参数 init
        println("3. 初始化流水线对象 填写参数 init")
        val myPipeline: MyPipeline = new MyPipeline().setLableColName("gender")
                .setFeatureColNames(Array("hair", "height", "skirt", "age"))
                .setMaxCategories(10).init()


        //4. 把数据送入流水线进行训练
        println("4. 把数据送入流水线进行训练")
        trainDF.show(1000,false)
        myPipeline.train(trainDF)

        //5. 把决策树进行打印 观察  特征权重
        println("5. 把决策树进行打印 观察  特征权重")
        myPipeline.printTree()
        myPipeline.printFeatureWeight()

        //6. 预测
        println("6. 预测")
        val predictedDF: DataFrame = myPipeline.predict(testDF)
        predictedDF.show(1000,false)
    }

}
